from sklearn.feature_selection import RFE, VarianceThreshold, mutual_info_classif, SelectKBest
from sklearn.ensemble import RandomForestClassifier  # 可换模型
from sklearn.preprocessing import StandardScaler

from utils.utils import *
class RFEUtil:
    def __init__(self,dp,split_random_state,pic_save_path,feature_save_path=rfe_feature_path):
        self.total_df = None
        self.feature_save_path = opj(feature_save_path,str(split_random_state))
        self.pic_save_path = opj(pic_save_path,str(split_random_state),'feature_extract','rfe')
        self.dp = dp


    def extract_feature(self,df,data_type,n_features_to_select=5):

        # 先划分训练/测试集
        X_train, X_test, y_train, y_test, _, _ = self.dp.split_train_and_test(df.copy())
        vt = VarianceThreshold(threshold=1e-5)
        X_train_vt = vt.fit_transform(X_train)
        vt_feature_names = X_train.columns[vt.get_support()]

        skb = SelectKBest(score_func=mutual_info_classif, k=('all' if X_train_vt.shape[1] < X_train_vt.shape[0] else 200))
        X_train_selected = skb.fit_transform(X_train_vt, y_train)
        skb_feature_names = vt_feature_names[skb.get_support()]

        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train_selected)
        X_train_scaled = scaler.fit_transform(X_train_selected)
        import pandas as pd
        X_train_scaled_df = pd.DataFrame(X_train_scaled, columns=skb_feature_names, index=X_train.index)


        model = RandomForestClassifier(n_estimators=100, random_state=42)
        selector = RFE(model, n_features_to_select=n_features_to_select)
        selector.fit(X_train_scaled_df, y_train)
        selected_features = X_train_scaled_df.columns[selector.get_support()].tolist()
        df_filtered = df.loc[:, ['Patient_ID', target_column] + selected_features]
        df_filtered = df_filtered.dropna(how='all', subset=selected_features    )

        md(self.feature_save_path)
        df_filtered.to_csv(opj(self.feature_save_path, data_type.replace('FeatureSummary', '') + '.csv'), index=False)
        if self.total_df is None:
            self.total_df = df_filtered
        else:
            self.total_df = self.total_df.merge(df_filtered, how='left', on=['Patient_ID', target_column])
        if len(selected_features) < 2:
            return True
        return False

    def save_total_df(self):
        md(self.feature_save_path)
        self.total_df.to_csv(opj(self.feature_save_path, 'total.csv'), index=False)